import mplfinance as mpf
import pandas as pd

# 获取数据，处理数据
def read_data(filename):
  df = pd.read_csv('./files/history_A_stock_k_data.csv')
  df['Open'] = df['open']
  df['Close'] = df['close']
  df['High'] = df['high']
  df['Low'] = df['low']
  df['Volume'] = df['volume']
  df['Date'] = df['date']

  df['Date'] = pd.to_datetime(df['Date'])
  df.set_index(['Date'], inplace=True)
  return df


# 获取数据
df = read_data('./files/test02.csv')

# 设置线元素的颜色
my_color = mpf.make_marketcolors(
  up="red",  # 上涨K线的颜色
  down="green",  # 下跌K线的颜色
  edge="black",  # 蜡烛图箱体的颜色
  volume="purple",  # 成交量柱子的颜色
  wick="black"  # 蜡烛图影线的颜色
)

# 自定义风格
my_style = mpf.make_mpf_style(
  base_mpf_style='blueskies',
  # base_mpl_style='seaborn',  # 也可以试试matplotlib的seaborn等风格。
  marketcolors=my_color,
  figcolor='(0, 0.8, 0.85)',
  gridcolor='(0.9, 0.9, 0.9)',
  rc={'font.family': 'SimHei', 'axes.unicode_minus': 'False'}
)

print(df.columns)
# 选择平安银行2023年7月11日到2024年5月17日的数据进行绘图
mpf.plot(
  df,
  type='candle',
  ylabel="open",
  style=my_style,
  title='sh600508 日线行情',
  mav=(5, 10),
  volume=True,
  figratio=(5, 3),
  ylabel_lower="vol", 
  # savefig='sh600000.jpg' #保存为图片
)

